import cv2
import time
import numpy as np
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from PIL import Image
from coppeliasim_zmqremoteapi_client import RemoteAPIClient
import os
import csv


def main():
    frame_count = 0
    zNear = 0.0100
    zFar = 4.0000
    file_path = "data/sim_matrice.csv"

    first_run = True

    print('Program started')

    if not os.path.exists('experiment_data'):
        os.makedirs('data')
    rgba_images_path = os.path.join('data', 'rgba_images')
    depth_images_path = os.path.join('data', 'depth_images')
    if not os.path.exists(rgba_images_path):
        os.makedirs(rgba_images_path)
    if not os.path.exists(depth_images_path):
        os.makedirs(depth_images_path)

    if os.path.exists(file_path):
        print(f"Removing existing file: {file_path}")
        os.remove(file_path)

    with open(file_path, 'w', newline='') as file:
        writer = csv.writer(file)
        # 写入表头
        writer.writerow(["Index", "Vx0", "Vy0", "Vz0", "P0", "Vx1", "Vy1", "Vz1", "P1", "Vx2", "Vy2", "Vz2", "P2"])

    client = RemoteAPIClient()
    sim = client.require('sim')

    visionSensorHandle = sim.getObject('/visionSensor')

    blockHandle = sim.getObject('/Cuboid')

    _, resolution = sim.getVisionSensorImg(visionSensorHandle)

    sim.setStepping(True)# 设置为步进模式
    sim.startSimulation()

    try:
        while True:
            if not first_run:
                frame_count += 1
            print(f'Frame: {frame_count}')

            img, res = sim.getVisionSensorImg(visionSensorHandle)
            depth_bytes,__ = sim.getVisionSensorDepth(visionSensorHandle)

            with open(file_path, 'a', newline='') as file:
                writer = csv.writer(file)
                # 获取变换矩阵
                matrix = sim.getObjectMatrix(blockHandle)
                # 写入文件，格式为：编号, 矩阵值
                if not first_run:
                    writer.writerow([frame_count] + matrix)

            sim.step()# 仿真进行一步
            
            display.displayUpdated(img, res)

            color_img = img = np.frombuffer(img, dtype=np.uint8).reshape(resolution[1], resolution[0], 3)
            color_img = color_img[ :,::-1, :]
            rgba_filename = f'data/rgba_images/frame_{frame_count:04d}.png'
            color_img = cv2.cvtColor(color_img, cv2.COLOR_BGR2RGBA)
            if not first_run:
                cv2.imwrite(rgba_filename, color_img)

            depth_img = sim.unpackFloatTable(depth_bytes)
            depth_img = np.array(depth_img).reshape((resolution[1], resolution[0]))
            depth_img = np.fliplr(depth_img)
            depth_img = depth_img * (zFar - zNear) + zNear
            
            depth_img = (depth_img * 1000).astype(np.uint16)
            depth_filename = f'data/depth_images/frame_{frame_count:04d}.png'
            if not first_run:
                cv2.imwrite(depth_filename, depth_img)
            first_run = False
    except KeyboardInterrupt:
        print('Program stopped by user.')
    finally:

        sim.stopSimulation()
        print('Program ended')

        

if __name__ == '__main__':
    main()